import checker as chk
from collections import deque

AvailableOp = []
AvailableCmp = []
AvailableBinaryLogic = []
AvailableUnaryLogic = []
AvailableConst = []
AvailableArg = []
ArgIdx = {}
BranchAvailable = False
EqualCmpAvailable = False
LessCmpAvailable = False
PlusAvailable = False
MinusAvailable = False
AllowSynthArith = False
AllowSynthConst = False

ArrayElementArgs = []
ArrayArgs = []
ArrayIndexOffset = []

ARRAY_IDX = "!!IDX"
ARRAY_PREV = "!!PRV"
ARRAY_CUR = "!!CUR"
SUPPORTED_OP = ['+', '-']
SUPPORTED_CMP = ['=', '>', '<', '>=', '<=']
SUPPORTED_BIN_LOGIC = ['and', 'or', '=>']
SUPPORTED_UN_LOGIC = ['not']

IsReduceTaskSynth = False
ReduceTaskSynthState = None
BQ = deque()
RQ = deque()
Ck = None
FunDefineStr = None
ReduceTaskLen = 0


class NotSimpleSynthesisError(Exception):
    pass


def solve_fail():
    raise NotSimpleSynthesisError()


def check_syntax_capability():
    global PlusAvailable
    global MinusAvailable
    global EqualCmpAvailable
    global LessCmpAvailable
    if '+' in AvailableOp:
        PlusAvailable = True
    if '-' in AvailableOp:
        MinusAvailable = True
    if '=' in AvailableCmp:
        EqualCmpAvailable = True
    for cmp in ['>', '<', '>=', '<=']:
        if cmp in AvailableCmp:
            LessCmpAvailable = True


def is_non_empty_list(x):
    return type(x) == list and len(x) > 0


def is_binary_op(x, nt):
    return type(x) == list and len(x) == 3 and type(x[0]) == str and x[1] == nt and x[2] == nt


def is_unary_op(x, nt):
    return type(x) == list and len(x) == 2 and type(x[0]) == str and x[1] == nt


def is_ite_stmt(x, ni, nb):
    return type(x) == list and len(x) == 4 and x[0] == 'ite' and x[1] == nb and x[2] == ni and x[3] == ni


def constraint_item_scan(stmt):
    global AllowSynthArith
    global AllowSynthConst
    if type(stmt) == tuple:
        if stmt[0] == 'Int':
            AllowSynthConst = True
    elif type(stmt) == list:
        for sub in stmt:
            constraint_item_scan(sub)
    else:
        if stmt in ['+', '-', '*', '/', 'div', 'mod']:
            AllowSynthArith = True


def constraint_scan(bm_expr):
    for x in bm_expr:
        if type(x) == list and len(x) > 0 and x[0] == 'constraint':
            constraint_item_scan(x)


##############################
# synthesis target components
##############################

class SItem:
    def __init__(self):
        self.has_nt = False
        self.nt_depth = -1
        self.sa = None

    def check_sub_nt(self, x):
        if x.has_nt and (x.nt_depth + 1 < self.nt_depth or not self.has_nt):
            self.nt_depth = x.nt_depth + 1
            self.has_nt = True

    def sexpr(self):
        raise NotImplementedError("abstract class")

    def evaluate(self, f_arg):
        raise NotImplementedError("abstract class")

    def expand(self, si):
        return []

    def get_arg_or_const(self):
        return None


class NTInt(SItem):
    def __init__(self):
        super().__init__()
        self.has_nt = True
        self.nt_depth = 0
        self.sa = set()

    def sexpr(self):
        return None

    def evaluate(self, f_arg):
        return None

    def expand(self, si):
        sr = []
        kc = list(filter(lambda i: type(i) == int, si[2]))
        kc = None if len(kc) == 0 else min(kc)
        ka = list(filter(lambda i: type(i) == str, si[2]))
        ka = None if len(ka) == 0 else min(ka)
        if IsReduceTaskSynth:
            for x in ArrayArgs:
                if x not in si[1] and (ka is None or x >= ka):
                    sr.append(IntArg(x))
        else:
            for x in AvailableArg:
                if x not in si[1] and kc is None and (ka is None or x >= ka):
                    sr.append(IntArg(x))
            if AllowSynthConst:
                for c in AvailableConst:
                    if c not in si[1] and (kc is None or c >= kc):
                        sr.append(IntConst(c))
        if AllowSynthArith:
            if PlusAvailable and '+' in si[0]:
                sr.append(Plus())
            if MinusAvailable and '-' in si[0]:
                sr.append(Minus())
        if BranchAvailable:
            if LessCmpAvailable and '<' in si[0]:
                sr.append(IfLessInt())
            if EqualCmpAvailable and '=' in si[0]:
                sr.append(IfEqualInt())
        return sr


class IntArg(SItem):
    def __init__(self, x: str):
        super().__init__()
        self.__x = x
        self.sa = {x}

    def sexpr(self):
        if IsReduceTaskSynth:
            if self.__x == ARRAY_PREV:
                return ReduceTaskSynthState[0]
            elif self.__x == ARRAY_CUR:
                return ReduceTaskSynthState[1]
            elif self.__x == ARRAY_IDX:
                return ReduceTaskSynthState[2]
        return self.__x

    def evaluate(self, f_arg):
        return f_arg[ArgIdx[self.__x]]

    def get_arg_or_const(self):
        return self.__x


class IntConst(SItem):
    def __init__(self, c: int):
        super().__init__()
        self.__c = c
        self.sa = {c}

    def sexpr(self):
        return str(self.__c)

    def evaluate(self, f_arg):
        return self.__c

    def get_arg_or_const(self):
        return self.__c


class Plus(SItem):
    def __init__(self, lhs: SItem = NTInt(), rhs: SItem = NTInt()):
        super().__init__()
        self.__lhs = lhs
        self.__rhs = rhs
        self.check_sub_nt(self.__lhs)
        self.check_sub_nt(self.__rhs)
        self.sa = self.__lhs.sa.copy()
        self.sa.update(self.__rhs.sa)

    def sexpr(self):
        return '(+ %s %s)' % (self.__lhs.sexpr(), self.__rhs.sexpr())

    def evaluate(self, f_arg):
        return self.__lhs.evaluate(f_arg) + self.__rhs.evaluate(f_arg)

    def expand(self, si):
        if 0 not in si[1]:
            si = (si[0], {0}, si[2])
        sd = si[2].copy()
        sd.update(self.__lhs.sa)
        sd.update(self.__rhs.sa)
        sr = []
        # right
        sub_sr = self.__rhs.expand(([], si[1], sd))
        for sub in sub_sr:
            sr.append(Plus(self.__lhs, sub))
        if len(sr) > 0:
            return sr
        # left
        sub_sr = self.__lhs.expand((['+'], si[1], sd))
        for sub in sub_sr:
            sr.append(Plus(sub, self.__rhs))
        return sr


class Minus(SItem):
    def __init__(self, lhs: SItem = NTInt(), rhs: SItem = NTInt()):
        super().__init__()
        self.__lhs = lhs
        self.__rhs = rhs
        self.check_sub_nt(self.__lhs)
        self.check_sub_nt(self.__rhs)
        self.sa = self.__lhs.sa.copy()
        self.sa.update(self.__rhs.sa)

    def sexpr(self):
        return '(- %s %s)' % (self.__lhs.sexpr(), self.__rhs.sexpr())

    def evaluate(self, f_arg):
        return self.__lhs.evaluate(f_arg) - self.__rhs.evaluate(f_arg)

    def expand(self, si):
        sr = []
        # left
        if self.__lhs.has_nt and self.__lhs.nt_depth == self.nt_depth - 1:
            da = {0}
            da.update(self.__rhs.sa)
            sub_sr = self.__lhs.expand((['+'], da, si[2]))
            for sub in sub_sr:
                sr.append(Minus(sub, self.__rhs))
            if len(sr) > 0:
                return sr
        # right
        da = {0}
        da.update(self.__lhs.sa)
        sub_sr = self.__rhs.expand((['+'], da, si[2]))
        for sub in sub_sr:
            sr.append(Minus(self.__lhs, sub))
        return sr


def get_cmp_sig(c: str, lhs: SItem, rhs: SItem):
    lc = lhs.get_arg_or_const()
    rc = rhs.get_arg_or_const()
    if lc is not None and rc is not None:
        return c, lc, rc
    return None


def rep_cmp_sig(c: str, oc):
    return c, oc[1], oc[2]


def rev_cmp_sig(oc):
    return oc[0], oc[2], oc[1]


def redundant_cmp(s: set, c):
    if c is not None:
        if c in s:
            return True
    return False


def same_const(a: SItem, b: SItem):
    ka = a.get_arg_or_const()
    kb = b.get_arg_or_const()
    return ka is not None and kb is not None and ka == kb


def unordered_const(a: SItem, b: SItem):
    ka = a.get_arg_or_const()
    kb = b.get_arg_or_const()
    if type(ka) == str:
        return type(kb) == str and ka >= kb
    elif type(ka) == int:
        return type(kb) == str or (type(kb) == int and ka >= kb)
    else:
        return False


class IfLessInt(SItem):
    def __init__(self, lc: SItem = NTInt(), rc: SItem = NTInt(), tr: SItem = NTInt(), fl: SItem = NTInt()):
        super().__init__()
        self.__lc = lc
        self.__rc = rc
        self.__tr = tr
        self.__fl = fl
        self.check_sub_nt(lc)
        self.check_sub_nt(rc)
        self.check_sub_nt(tr)
        self.check_sub_nt(fl)

    def sexpr(self):
        lcs = self.__lc.sexpr()
        rcs = self.__rc.sexpr()
        trs = self.__tr.sexpr()
        fls = self.__fl.sexpr()
        if '<' in AvailableCmp:
            return "(ite (< %s %s) %s %s)" % (lcs, rcs, trs, fls)
        elif '>' in AvailableCmp:
            return "(ite (> %s %s) %s %s)" % (rcs, lcs, trs, fls)
        elif '>=' in AvailableCmp:
            return "(ite (>= %s %s) %s %s)" % (lcs, rcs, fls, trs)
        elif '<=' in AvailableCmp:
            return "(ite (<= %s %s) %s %s)" % (rcs, lcs, fls, trs)
        else:
            return None

    def evaluate(self, f_arg):
        if self.__lc.evaluate(f_arg) < self.__rc.evaluate(f_arg):
            return self.__tr.evaluate(f_arg)
        else:
            return self.__fl.evaluate(f_arg)

    def expand(self, si):
        sr = []
        # LC
        if type(self.__lc) == NTInt:
            sub_sr = self.__lc.expand(([], si[1], si[2]))
            for sub in sub_sr:
                sr.append(IfLessInt(sub, self.__rc, self.__tr, self.__fl))
            return sr
        # RC
        if type(self.__rc) == NTInt:
            sub_sr = self.__rc.expand(([], si[1], si[2]))
            for sub in sub_sr:
                kc = get_cmp_sig('<', self.__lc, sub)
                if type(self.__lc) == IntConst and type(sub) == IntConst:
                    continue
                if same_const(self.__lc, sub):
                    continue
                if redundant_cmp(si[1], kc):
                    continue
                sr.append(IfLessInt(self.__lc, sub, self.__tr, self.__fl))
            return sr
        # THEN
        ka = get_cmp_sig('<', self.__lc, self.__rc)
        da = si[1].copy()
        da.add(ka)
        da.add(rep_cmp_sig('=', ka))
        da.add(rev_cmp_sig(ka))
        if self.__tr.has_nt and self.__tr.nt_depth + 1 == self.nt_depth:
            sub_sr = self.__tr.expand((['<', '=', '+', '-'], da, si[2]))
            for sub in sub_sr:
                if same_const(sub, self.__fl):
                    continue
                sr.append(IfLessInt(self.__lc, self.__rc, sub, self.__fl))
            if len(sr) > 0:
                return sr
        # ELSE
        da = si[1].copy()
        da.add(ka)
        if EqualCmpAvailable:
            da.add(rev_cmp_sig(ka))
        sub_sr = self.__fl.expand((['<', '=', '+', '-'], da, si[2]))
        for sub in sub_sr:
            if same_const(self.__tr, sub):
                continue
            sr.append(IfLessInt(self.__lc, self.__rc, self.__tr, sub))
        return sr


class IfEqualInt(SItem):
    def __init__(self, lc: SItem = NTInt(), rc: SItem = NTInt(), tr: SItem = NTInt(), fl: SItem = NTInt()):
        super().__init__()
        self.__lc = lc
        self.__rc = rc
        self.__tr = tr
        self.__fl = fl
        self.check_sub_nt(lc)
        self.check_sub_nt(rc)
        self.check_sub_nt(tr)
        self.check_sub_nt(fl)

    def sexpr(self):
        lcs = self.__lc.sexpr()
        rcs = self.__rc.sexpr()
        trs = self.__tr.sexpr()
        fls = self.__fl.sexpr()
        return "(ite (= %s %s) %s %s)" % (lcs, rcs, trs, fls)

    def evaluate(self, f_arg):
        if self.__lc.evaluate(f_arg) == self.__rc.evaluate(f_arg):
            return self.__tr.evaluate(f_arg)
        else:
            return self.__fl.evaluate(f_arg)

    def expand(self, si):
        sr = []
        # LC
        if type(self.__lc) == NTInt:
            sub_sr = self.__lc.expand(([], si[1], si[2]))
            for sub in sub_sr:
                sr.append(IfEqualInt(sub, self.__rc, self.__tr, self.__fl))
            return sr
        # RC
        if type(self.__rc) == NTInt:
            sub_sr = self.__rc.expand(([], si[1], si[2]))
            for sub in sub_sr:
                kc = get_cmp_sig('=', self.__lc, sub)
                if type(self.__lc) == IntConst and type(sub) == IntConst:
                    continue
                if unordered_const(self.__lc, sub):
                    continue
                if redundant_cmp(si[1], kc) or redundant_cmp(si[1], rev_cmp_sig(kc)):
                    continue
                sr.append(IfEqualInt(self.__lc, sub, self.__tr, self.__fl))
            return sr
        # THEN
        ka = get_cmp_sig('=', self.__lc, self.__rc)
        da = si[1].copy()
        da.add(ka)
        da.add(rep_cmp_sig('<', ka))
        da.add(rep_cmp_sig('<', rev_cmp_sig(ka)))
        if self.__tr.has_nt and self.__tr.nt_depth + 1 == self.nt_depth:
            sub_sr = self.__tr.expand((['=', '+', '-'], da, si[2]))
            for sub in sub_sr:
                if same_const(sub, self.__fl):
                    continue
                sr.append(IfEqualInt(self.__lc, self.__rc, sub, self.__fl))
            if len(sr) > 0:
                return sr
        # ELSE
        da = si[1].copy()
        da.add(ka)
        sub_sr = self.__fl.expand((['=', '+', '-'], da, si[2]))
        for sub in sub_sr:
            if same_const(self.__tr, sub):
                continue
            sr.append(IfEqualInt(self.__lc, self.__rc, self.__tr, sub))
        return sr


##############################
# reduce task detection
##############################

def argument_sanity_check(arg_names):
    for a in arg_names:
        if a not in AvailableArg:
            return False
    for a in AvailableArg:
        if a not in arg_names:
            return False
    return True


def detect_indexed_search_task(a_arg, x_arg):
    # Precondition: len(a_arg)>=3
    arg_prefix = a_arg[0][:-1]
    start_idx = 0 if a_arg[0][-1] == '0' else 1
    for i in range(len(a_arg)):
        expected_name = arg_prefix + str(i + start_idx)
        if expected_name != a_arg[i]:
            return 0
    task_len = len(a_arg)
    ok_0 = True
    ok_1 = True
    for i in range(task_len):
        if i not in AvailableConst:
            ok_0 = False
        if (i + 1) not in AvailableConst:
            ok_1 = False
    if not (ok_0 or ok_1):
        return 0
    # success
    if ok_0:
        ArrayIndexOffset.append(0)
    if ok_1:
        ArrayIndexOffset.append(1)
    if x_arg is not None:
        ArrayArgs.append(x_arg)
    ArrayArgs.append(ARRAY_IDX)
    ArrayArgs.append(ARRAY_PREV)
    ArrayArgs.append(ARRAY_CUR)
    for a in a_arg:
        ArrayElementArgs.append(a)
    return task_len


def detect_ordinary_reduce_task(a_arg):
    # Precondition: len(a_arg)>3
    task_len = len(a_arg)
    ArrayArgs.append(ARRAY_PREV)
    ArrayArgs.append(ARRAY_CUR)
    for a in a_arg:
        ArrayElementArgs.append(a)
    return task_len


def get_full_reduce_func(idx_offset, initial_value, reduce_task_len, cur):
    global ReduceTaskSynthState
    last_code = str(initial_value) if initial_value is not None else ArrayElementArgs[0]
    start_idx = 0 if initial_value is not None else 1
    for i in range(start_idx, reduce_task_len):
        ReduceTaskSynthState = (last_code, ArrayElementArgs[i], str(i + idx_offset))
        last_code = cur.sexpr()
    return last_code


##############################
# main
##############################

def enhanced_synth_prepare(bm_expr):
    global BranchAvailable
    global IsReduceTaskSynth
    global ReduceTaskSynthState
    global FunDefineStr
    global ReduceTaskLen
    global Ck
    synth_fun_expr = None
    for expr in bm_expr:
        if is_non_empty_list(expr) and expr[0] == 'synth-fun':
            synth_fun_expr = expr
    fun_define = ['define-fun'] + synth_fun_expr[1:4]
    FunDefineStr = chk.list_to_string(fun_define)
    bool_nt_name = None
    int_nt_name = None
    bool_nt_cnt = 0
    int_nt_cnt = 0

    fun_args = synth_fun_expr[2]
    fun_arg_names = list([x[0] for x in fun_args])
    for i in range(len(fun_args)):
        if fun_args[i][1] != 'Int':
            solve_fail()
        ArgIdx[fun_args[i][0]] = i

    if synth_fun_expr[3] != 'Int':
        solve_fail()

    for nt_item in synth_fun_expr[4]:
        nt_type = nt_item[1]
        if nt_type == 'Bool':
            bool_nt_name = nt_item[0]
            bool_nt_cnt += 1
        elif nt_type == 'Int':
            int_nt_name = nt_item[0]
            int_nt_cnt += 1
        else:
            solve_fail()

    if bool_nt_cnt > 1 or int_nt_cnt != 1:
        solve_fail()

    for nt_item in synth_fun_expr[4]:
        if nt_item[0] == int_nt_name:
            for ntt in nt_item[2]:
                if type(ntt) == tuple:
                    if ntt[0] != 'Int':
                        solve_fail()
                    AvailableConst.append(ntt[1])
                elif type(ntt) == str:
                    AvailableArg.append(ntt)
                elif is_binary_op(ntt, int_nt_name) and ntt[0] in SUPPORTED_OP:
                    AvailableOp.append(ntt[0])
                elif is_ite_stmt(ntt, int_nt_name, bool_nt_name):
                    BranchAvailable = True
                else:
                    solve_fail()
        elif nt_item[0] == bool_nt_name:
            for ntt in nt_item[2]:
                if is_binary_op(ntt, int_nt_name) and ntt[0] in SUPPORTED_CMP:
                    AvailableCmp.append(ntt[0])
                elif is_binary_op(ntt, bool_nt_name) and ntt[0] in SUPPORTED_BIN_LOGIC:
                    AvailableBinaryLogic.append(ntt[0])
                elif is_unary_op(ntt, bool_nt_name) and ntt[0] in SUPPORTED_UN_LOGIC:
                    AvailableUnaryLogic.append(ntt[0])
                else:
                    solve_fail()
    check_syntax_capability()
    constraint_scan(bm_expr)
    # print(int_nt_name)
    # print(bool_nt_name)
    # print(AvailableOp)
    # print(AvailableCmp)
    # print(AvailableArg)
    # print(AvailableConst)
    # print(AvailableBinaryLogic)
    # print(AvailableUnaryLogic)
    # print(BranchAvailable)
    # print(ArgIdx)
    # print(str([EqualCmpAvailable, LessCmpAvailable, PlusAvailable, MinusAvailable]))
    ReduceTaskLen = 0
    if argument_sanity_check(fun_arg_names) and len(fun_arg_names) > 3:
        ReduceTaskLen = detect_indexed_search_task(fun_arg_names, None)
        if ReduceTaskLen == 0:
            ReduceTaskLen = detect_indexed_search_task(fun_arg_names[:-1], fun_arg_names[-1])
        if ReduceTaskLen == 0:
            ReduceTaskLen = detect_indexed_search_task(fun_arg_names[1:], fun_arg_names[0])
        if ReduceTaskLen == 0:
            ReduceTaskLen = detect_ordinary_reduce_task(fun_arg_names)
    Ck = chk.Checker(bm_expr)
    BQ.append(NTInt())
    if ReduceTaskLen > 0:
        RQ.append(NTInt())
        if ARRAY_IDX not in ArrayArgs:
            ArrayIndexOffset.append(-999999)
    # print(ReduceTaskLen)
    # print(ArrayElementArgs)
    # print(ArrayArgs)
    # print(ArrayIndexOffset)
    # raise NotImplementedError()


def normal_iteration():
    global IsReduceTaskSynth
    IsReduceTaskSynth = False
    if len(BQ) == 0:
        return None
    cur = BQ.popleft()
    cd = cur.nt_depth if cur.has_nt else -1
    if not cur.has_nt:
        cur_str = cur.sexpr()
        full_str = FunDefineStr[:-1] + ' ' + cur_str + FunDefineStr[-1]
        ce = Ck.get_counterexample(full_str)
        if ce is None:
            return full_str
    else:
        try_extend = cur.expand((['<', '=', '+', '-'], set(), set()))
        for te in try_extend:
            if not te.has_nt or te.nt_depth < cd:
                BQ.appendleft(te)
            else:
                BQ.append(te)
    return None


def reduce_task_iteration():
    global IsReduceTaskSynth
    IsReduceTaskSynth = True
    if len(RQ) == 0:
        return None
    cur = RQ.popleft()
    cd = cur.nt_depth if cur.has_nt else -1
    if not cur.has_nt:
        for idx_offset in ArrayIndexOffset:
            last_code = get_full_reduce_func(idx_offset, None, ReduceTaskLen, cur)
            full_str = FunDefineStr[:-1] + ' ' + last_code + FunDefineStr[-1]
            ce = Ck.get_counterexample(full_str)
            if ce is None:
                return full_str
            for initial_value in AvailableConst:
                last_code = get_full_reduce_func(idx_offset, initial_value, ReduceTaskLen, cur)
                full_str = FunDefineStr[:-1] + ' ' + last_code + FunDefineStr[-1]
                ce = Ck.get_counterexample(full_str)
                if ce is None:
                    return full_str
    else:
        try_extend = cur.expand((['<', '=', '+', '-'], set(), set()))
        for te in try_extend:
            if not te.has_nt or te.nt_depth < cd:
                RQ.appendleft(te)
            else:
                RQ.append(te)
    return None
